"""
https://www.cnblogs.com/long5683/p/9692845.html

扑克牌分水岭
"""
import cv2 as cv
import numpy as np
import os
import sys
import matplotlib.pyplot as plt

np.random.seed(1)

path = 'images/poker_cards.png'
src = cv.imread(path, cv.IMREAD_COLOR)
src_ = src.copy()

cv.imshow('1 input', src)

# 白色背景变成黑色
idx = src[:, :] == (255, 255, 255)
src[idx] = 0
cv.imshow('2 255->0', src)

# sharpen(提高对比度)
kernel = np.float32([
    [0, 1, 0],
    [1, -8, 1],
    [0, 1, 0]
])
# 拉普拉斯算子实现边缘提取
imgLaplance = cv.filter2D(src, cv.CV_32F, kernel, borderType=cv.BORDER_DEFAULT)  # 拉普拉斯有浮点数计算，位数要提高到32
# check and draw imgLaplance
imgLaplance_ = imgLaplance.copy()
imgLaplance_ = cv.normalize(imgLaplance_, None, 0., 255., norm_type=cv.NORM_MINMAX)
imgLaplance_ = imgLaplance_.astype(np.uint8)
cv.imshow('3 Laplance', imgLaplance_)

# 原图减边缘（白色）实现边缘增强
sharpenImg = src.astype(np.float32)
resultImg = sharpenImg - imgLaplance
# check and draw resultImg
resultImg_ = resultImg.copy()
cv.normalize(resultImg_, resultImg_, 0., 255., norm_type=cv.NORM_MINMAX)
resultImg_ = resultImg_.astype(np.uint8)
cv.imshow('4 subtract', resultImg_)

# 转换成二值图
resultImg = cv.cvtColor(resultImg_, cv.COLOR_BGR2GRAY)
ret, binary = cv.threshold(resultImg, 40, 255, cv.THRESH_OTSU + cv.THRESH_BINARY)
cv.imshow('5 binary', binary)

# 距离变换
distImg = cv.distanceTransform(binary, cv.DIST_L2, 3, dstType=cv.CV_32F)
cv.normalize(distImg, distImg, 0., 1., cv.NORM_MINMAX)
cv.imshow('6 distance', distImg)

# 二值化
ret, distImg = cv.threshold(distImg, 0.4, 1.0, cv.THRESH_BINARY)
cv.imshow('7 bin', distImg)

# 腐蚀(使得连在一起的部分分开)
k1 = np.ones([3, 3], dtype=np.uint8)
distImg = cv.erode(distImg, k1)
cv.imshow('7 erode', distImg)

# 标记
distImg *= 255.
dist_8u = distImg.astype(np.uint8)
contours, hierarchy = cv.findContours(dist_8u, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE, offset=(0, 0))
bg = np.zeros(src.shape, dtype=np.uint8)
cv.drawContours(bg, contours, -1, (0, 255, 0), 1)
cv.imshow('8 contours', bg)
# 创建标记
marker = np.zeros(distImg.shape, dtype=np.int32)
# 画标记
for i, c in enumerate(contours):
    cv.drawContours(marker, [c], 0, i + 1, -1)
# 可视化标记
# color table start
# 随机生成几种颜色
colorTab = [(255, 255, 255)]
for i in range(len(contours)):
    b = np.random.randint(0, 256)
    g = np.random.randint(0, 256)
    r = np.random.randint(0, 256)
    colorTab.append((b, g, r))
colorTab = np.uint8(colorTab)
# color table end
print(np.unique(marker))
marker_drawing = colorTab[marker]
cv.imshow('8 marker', marker_drawing)

# 分水岭变换
src = src_.copy()
cv.watershed(src, marker)
print(np.unique(marker))
marker_drawing = marker
marker_drawing[marker_drawing == -1] = 0
print(np.unique(marker_drawing))
marker_drawing = colorTab[marker]
cv.imshow('9 after water', marker_drawing)

cv.waitKey(0)
